# coding:utf-8
#
# The MIT License (MIT)
#
# Copyright (c) 2016-2018 yutiansut/QUANTAXIS
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
#


from functools import reduce
import math
import numpy as np
import pandas as pd

"""
使用方法：import series_tai as stai
Series 类
# 第一个参数无特别说明，类型为 pd.Series
# 返回类型为 pd.Series
这个是下面以DataFrame为输入的基础函数
return new pd.Series 
"""


def EMA(Series, N):
    return pd.Series.ewm(Series, span=N, min_periods=1, adjust=True).mean()


def MA(Series, N):
    return pd.Series.rolling(Series, N, min_periods=1).mean()


# 威廉SMA  参考https://www.joinquant.com/post/867


def SMA(Series, N, M=1):
    """
    威廉SMA算法

    本次修正主要是对于返回值的优化,现在的返回值会带上原先输入的索引index
    2018/5/3
    @yutiansut
    """
    ret = []
    i = 1
    length = len(Series)
    # 跳过X中前面几个 nan 值
    while i < length:
        if np.isnan(Series.iloc[i]):
            i += 1
        else:
            break
    preY = Series.iloc[i]  # Y'
    ret.append(preY)
    while i < length:
        Y = (M * Series.iloc[i] + (N - M) * preY) / float(N)
        ret.append(Y)
        preY = Y
        i += 1
    return pd.Series(ret, index=Series.tail(len(ret)).index)


def DIFF(Series, N=1):
    return pd.Series(Series).diff(N)


def HHV(Series, N):
    return pd.Series(Series).rolling(N).max()


def LLV(Series, N):
    return pd.Series(Series).rolling(N).min()


def SUM(Series, N):
    return pd.Series.rolling(Series, N).sum()


def ABS(Series):
    return abs(Series)


def MAX(A, B):
    var = IF(A > B, A, B)
    return var


def MIN(A, B):
    var = IF(A < B, A, B)
    return var


def SINGLE_CROSS(A, B):
    if A.iloc[-2] < B.iloc[-2] and A.iloc[-1] > B.iloc[-1]:
        return True
    else:
        return False


def CROSS(A, B):
    """A<B then A>B  A上穿B B下穿A

    Arguments:
        A: pd.Series
        B: Pd.Series

    Returns:
       pd.Series, where 1==a up_cross b, -1 == a down_cross b
    """

    var = np.where(A > B, 1, 0)
    crss = pd.Series(var, index=A.index).diff(1)
    crss[0] = 0
    return crss  # .apply(int)


def COUNT(COND, N):
    """
    2018/05/23 修改

    参考https://github.com/QUANTAXIS/QUANTAXIS/issues/429

    现在返回的是series
    """
    return pd.Series(np.where(COND, 1, 0), index=COND.index).rolling(N).sum()


def IF(COND, V1, V2):
    var = np.where(COND, V1, V2)
    return pd.Series(var, index=V1.index)


def IFAND(COND1, COND2, V1, V2):
    var = np.where(np.logical_and(COND1, COND2), V1, V2)
    return pd.Series(var, index=V1.index)


def IFOR(COND1, COND2, V1, V2):
    var = np.where(np.logical_or(COND1, COND2), V1, V2)
    return pd.Series(var, index=V1.index)


def REF(Series, N):
    var = Series.diff(N)
    var = Series - var
    return var


def LAST(COND, N1, N2):
    """表达持续性
    从前N1日到前N2日一直满足COND条件

    Arguments:
        COND {[type]} -- [description]
        N1 {[type]} -- [description]
        N2 {[type]} -- [description]
    """
    N2 = 1 if N2 == 0 else N2
    assert N2 > 0
    assert N1 > N2
    return COND.iloc[-N1:-N2].all()


def STD(Series, N):
    return pd.Series.rolling(Series, N).std()


def AVEDEV(Series, N):
    """
    平均绝对偏差 mean absolute deviation
    修正: 2018-05-25

    之前用mad的计算模式依然返回的是单值
    """
    return Series.rolling(N).apply(lambda x: (np.abs(x - x.mean())).mean(), raw=True)


def MACD(Series, FAST, SLOW, MID):
    """macd指标 仅适用于Series
    对于DATAFRAME的应用请使用QA_indicator_macd
    """
    EMAFAST = EMA(Series, FAST)
    EMASLOW = EMA(Series, SLOW)
    DIFF = EMAFAST - EMASLOW
    DEA = EMA(DIFF, MID)
    MACD = (DIFF - DEA) * 2
    DICT = {'DIFF': DIFF, 'DEA': DEA, 'MACD': MACD}
    VAR = pd.DataFrame(DICT)
    return VAR


def BBIBOLL(Series, N1, N2, N3, N4, N, M):  # 多空布林线

    bbiboll = BBI(Series, N1, N2, N3, N4)
    UPER = bbiboll + M * STD(bbiboll, N)
    DOWN = bbiboll - M * STD(bbiboll, N)
    DICT = {'BBIBOLL': bbiboll, 'UPER': UPER, 'DOWN': DOWN}
    VAR = pd.DataFrame(DICT)
    return VAR


def BBI(Series, N1, N2, N3, N4):
    '多空指标'

    bbi = (MA(Series, N1) + MA(Series, N2) +
           MA(Series, N3) + MA(Series, N4)) / 4
    DICT = {'BBI': bbi}
    VAR = pd.DataFrame(DICT)
    return VAR


def TR(df, sgn_close='close', sgn_high='high', sgn_low='low', n=5):
    """
    计算真实波幅
    【输入】
        df, pd.DataFrame 格式数据源
        sgn_close，列名，一般是：close收盘价
    【输出】
        pd.DataFrame格式数据源,

    """

    i = 0
    TR_l = [df[sgn_high][i] - df[sgn_low][0]]
    while i < len(df) - 1:  # df.index[-1]:
        TR = max(df[sgn_high].iloc[i + 1], df[sgn_close].iloc[i]) - min(df[sgn_low].iloc[i + 1], df[sgn_close].iloc[i])
        TR_l.append(TR)
        i += 1
    TR_s = pd.Series(TR_l, index=df.index, name='TR')
    ATR = pd.Series(TR_s.ewm(span=n, min_periods=n).mean(), name='ATR', index=df.index)
    return TR_s, ATR


def ACSP(Series, spp=5, acp=5):
    """速度加速度
    速度：价格的离差
    加速度：速度的离差"""
    sp = Series.diff(spp)
    ac = sp.diff(acp)
    return sp, ac


def BARLAST(cond, yes=True):
    """支持MultiIndex的cond和DateTimeIndex的cond
    条件成立  yes= True 或者 yes=1 根据不同的指标自己定

    Arguments:
        cond {[type]} -- [description]
    """
    if isinstance(cond.index, pd.MultiIndex):
        return len(cond) - cond.index.levels[0].tolist().index(cond[cond != yes].index[-1][0]) - 1
    elif isinstance(cond.index, pd.DatetimeIndex):
        return len(cond) - cond.index.tolist().index(cond[cond != yes].index[-1]) - 1


ma = MA
ema = EMA
sma = SMA
diff = DIFF
hhv = HHV
llv = LLV
sum = SUM
cross = CROSS
single_cross = SINGLE_CROSS
count = COUNT
ifand = IFAND
ifor = IFOR
ref = REF
last = LAST
std = STD
avedev = AVEDEV
macd = MACD
bbiboll = BBIBOLL
bbi = BBI
barlast = BARLAST

XARROUND = lambda x, y: np.round(y * (round(x / y - math.floor(x / y) + 0.00000000001) + math.floor(x / y)), 2)

if __name__ == '__main__':
    a = pd.Series([1, 6, 3, 1, 5, 2])
    b = pd.Series([5, 4, 5, 2, 1, 6])
    ma2 = ma(a, 4)
    print(f'ma={ma2}')
    c = CROSS(a, b)
    print(c < 0)

    print('a cross b :', c)
    print(' b cross a: ', cross(b, a))
    print(a - b)
    print(REF(a, 2))
    print(XARROUND(9, 5))
